'''
# 设计3D模型用于训练
# 简单3D直连式模型
# 3D-Unet改
# 3D-resnet50改
'''

import os
import time
import random
import shutil
import numpy as np
import pandas as pd
from tensorflow.keras.layers import Input, Conv3D, AveragePooling3D, MaxPooling3D, Flatten, Dense, UpSampling3D
from tensorflow.keras.layers import Activation, BatchNormalization, Dropout, Concatenate, Add
from tensorflow.keras.initializers import random_normal, glorot_uniform
from tensorflow.keras.models import Model, load_model
from tensorflow.keras.optimizers import SGD
from tensorflow.keras.utils import to_categorical, multi_gpu_model
from tensorflow.keras.callbacks import TensorBoard, ModelCheckpoint
from sklearn.metrics import classification_report, confusion_matrix

model_name = "five_3D"
time_str = time.strftime("%Y%m%d_%H%M%S", time.localtime())
log_dir = "/home/zmy/workspace2021/five_logs/" + model_name + time_str
work_dir = "/home/zmy/workspace2021/five_workdirs/" + model_name + time_str
model_save_path = "/home/zmy/workspace2021/five_models/" + model_name


def EasyNet(input_shape=(512, 512, 8, 2)):

    input = Input(input_shape, name='input_1')

    # (512,512,8)-->(256,256,4)
    X = AveragePooling3D(pool_size=(2, 2, 2), strides=(2, 2, 2), padding="valid", name='avgpool1')(input)

    # (256,256,4)-->(128,128,2)
    X = Conv3D(filters=64, kernel_size=(3, 3, 3), strides=(1, 1, 1), padding='same', name='conv1')(X)
    X = BatchNormalization(axis=4)(X)
    X = Activation('relu')(X)
    X = MaxPooling3D(pool_size=(2, 2, 2), strides=(2, 2, 2), padding="valid", name='maxpool1')(X)

    #  (128,128,2)-->(64,64,1)
    X = Conv3D(filters=64, kernel_size=(3, 3, 3), strides=(1, 1, 1), padding='same', name='conv2')(X)
    X = BatchNormalization(axis=4)(X)
    X = Activation('relu')(X)
    X = MaxPooling3D(pool_size=(2, 2, 2), strides=(2, 2, 2), padding="valid", name='maxpool2')(X)

    # (64,64,1)-->(32,32,1)
    X = Conv3D(filters=64, kernel_size=(3, 3, 3), strides=(1, 1, 1), padding='same', name='conv3')(X)
    X = BatchNormalization(axis=4)(X)
    X = Activation('relu')(X)
    X = MaxPooling3D(pool_size=(2, 2, 1), strides=(2, 2, 1), padding="valid", name='maxpool3')(X)

    # (32,32,1)-->(16,16,1)
    X = Conv3D(filters=64, kernel_size=(3, 3, 3), strides=(1, 1, 1), padding='same', name='conv4')(X)
    X = BatchNormalization(axis=4)(X)
    X = Activation('relu')(X)
    X = MaxPooling3D(pool_size=(2, 2, 1), strides=(2, 2, 1), padding="valid", name='maxpool4')(X)

    # (16,16,1)-->(8,8,1)
    X = Conv3D(filters=64, kernel_size=(3, 3, 3), strides=(1, 1, 1), padding='same', name='conv5')(X)
    X = BatchNormalization(axis=4)(X)
    X = Activation('relu')(X)
    X = MaxPooling3D(pool_size=(2, 2, 1), strides=(2, 2, 1), padding="valid", name='maxpool5')(X)

    # (8,8,1)-->(64,)
    X = Flatten()(X)
    output = Dense(5, activation='softmax', name='outclass', kernel_initializer=random_normal(seed=0))(X)

    # Create model
    model = Model(inputs=input, outputs=output, name='easy_model')
    return model


def Unet(input_shape=(512, 512, 8, 2)):
    input = Input(input_shape, name='input_1')

    # block1
    x1 = Conv3D(filters=32, kernel_size=(3, 3, 3), strides=(1, 1, 1), padding='same', name='conv11')(input)
    x1 = BatchNormalization(axis=1)(x1)
    x1 = Activation('relu')(x1)
    x1 = Conv3D(filters=64, kernel_size=(3, 3, 3), strides=(1, 1, 1), padding='same', name='conv12')(x1)
    x1 = BatchNormalization(axis=1)(x1)
    x1 = Activation('relu')(x1)


    # block2
    # (512,512,8,64)-->(128,128,4,128)
    x2 = MaxPooling3D(pool_size=(2, 2, 2), strides=(2, 2, 2), padding="valid", name='maxpool1')(x1)
    x2 = Conv3D(filters=64, kernel_size=(3, 3, 3), strides=(1, 1, 1), padding='same', name='conv21')(x2)
    x2 = BatchNormalization(axis=1)(x2)
    x2 = Activation('relu')(x2)
    x2 = Conv3D(filters=128, kernel_size=(3, 3, 3), strides=(1, 1, 1), padding='same', name='conv22')(x2)
    x2 = BatchNormalization(axis=1)(x2)
    x2 = Activation('relu')(x2)

    # block3
    # (128,128,4,128)-->(64,64,2,256)
    x3 = MaxPooling3D(pool_size=(2, 2, 2), strides=(2, 2, 2), padding="valid", name='maxpool2')(x2)
    x3 = Conv3D(filters=128, kernel_size=(3, 3, 3), strides=(1, 1, 1), padding='same', name='conv31')(x3)
    x3 = BatchNormalization(axis=1)(x3)
    x3 = Activation('relu')(x3)
    x3 = Conv3D(filters=256, kernel_size=(3, 3, 3), strides=(1, 1, 1), padding='same', name='conv32')(x3)
    x3 = BatchNormalization(axis=1)(x3)
    x3 = Activation('relu')(x3)

    # block4
    # (64,64,2,256)-->(32,32,1,512)
    x4 = MaxPooling3D(pool_size=(2, 2, 2), strides=(2, 2, 2), padding="valid", name='maxpool3')(x3)
    x4 = Conv3D(filters=256, kernel_size=(3, 3, 3), strides=(1, 1, 1), padding='same', name='conv41')(x4)
    x4 = BatchNormalization(axis=1)(x4)
    x4 = Activation('relu')(x4)
    x4 = Conv3D(filters=512, kernel_size=(3, 3, 3), strides=(1, 1, 1), padding='same', name='conv42')(x4)
    x4 = BatchNormalization(axis=1)(x4)
    x4 = Activation('relu')(x4)


    # up-block1
    x5 = UpSampling3D(size=(2,2,2))(x4)
    x5 = Concatenate(axis=4)([x3, x5])
    x5 = Conv3D(filters=256, kernel_size=(3, 3, 3), strides=(1, 1, 1), padding='same', name='convup11')(x5)
    x5 = BatchNormalization(axis=1)(x5)
    x5 = Activation('relu')(x5)
    x5 = Conv3D(filters=256, kernel_size=(3, 3, 3), strides=(1, 1, 1), padding='same', name='convup12')(x5)
    x5 = BatchNormalization(axis=1)(x5)
    x5 = Activation('relu')(x5)


    # up-block2
    x6 = UpSampling3D(size=(2, 2, 2))(x5)
    x6 = Concatenate(axis=4)([x2, x6])
    x6 = Conv3D(filters=128, kernel_size=(3, 3, 3), strides=(1, 1, 1), padding='same', name='convup21')(x6)
    x6 = BatchNormalization(axis=1)(x6)
    x6 = Activation('relu')(x6)
    x6 = Conv3D(filters=128, kernel_size=(3, 3, 3), strides=(1, 1, 1), padding='same', name='convup22')(x6)
    x6 = BatchNormalization(axis=1)(x6)
    x6 = Activation('relu')(x6)


    # up-block3
    x7 = UpSampling3D(size=(2, 2, 2))(x6)
    x7 = Concatenate(axis=4)([x1, x7])
    x7 = Conv3D(filters=64, kernel_size=(3, 3, 3), strides=(1, 1, 1), padding='same', name='convup31')(x7)
    x7 = BatchNormalization()(x7)
    x7 = Activation('relu')(x7)
    x7 = Conv3D(filters=64, kernel_size=(3, 3, 3), strides=(1, 1, 1), padding='same', name='convup32')(x7)
    x7 = BatchNormalization()(x7)
    x7 = Activation('relu')(x7)

    output = Conv3D(filters=5, kernel_size=(1, 1, 1), strides=(1, 1, 1), padding='same', name='outclass')(x7)
    print(output.shape)

    # Create model
    model = Model(inputs=input, outputs=output, name='easy_model')
    return model

def identity_block(X, f, filters, stage, block):
    """
    Implementation of the identity block as defined in Figure 3

    Arguments:
    X -- input tensor of shape (m, n_H_prev, n_W_prev, n_C_prev)
    f -- integer, specifying the shape of the middle CONV's window for the main path
    filters -- python list of integers, defining the number of filters in the CONV layers of the main path
    stage -- integer, used to name the layers, depending on their position in the network
    block -- string/character, used to name the layers, depending on their position in the network

    Returns:
    X -- output of the identity block, tensor of shape (n_H, n_W, n_C)
    """

    # defining name basis
    conv_name_base = 'res' + str(stage) + block + '_branch'
    bn_name_base = 'bn' + str(stage) + block + '_branch'

    # Retrieve Filters
    F1, F2, F3 = filters

    # Save the input value. You'll need this later to add back to the main path.
    X_shortcut = X

    # First component of main path
    X = Conv3D(filters=F1, kernel_size=(1, 1, 1), strides=(1, 1, 1), padding='same', name=conv_name_base + '2a',
               kernel_initializer=glorot_uniform(seed=0))(X)
    X = BatchNormalization(axis=4, name=bn_name_base + '2a')(X)
    X = Activation('relu')(X)

    # Second component of main path (≈3 lines)
    X = Conv3D(filters=F2, kernel_size=(f, f, f), strides=(1, 1, 1), padding='same', name=conv_name_base + '2b',
               kernel_initializer=glorot_uniform(seed=0))(X)
    X = BatchNormalization(axis=4, name=bn_name_base + '2b')(X)
    X = Activation('relu')(X)

    # Third component of main path (≈2 lines)
    X = Conv3D(filters=F3, kernel_size=(1, 1, 1), strides=(1, 1, 1), padding='same', name=conv_name_base + '2c',
               kernel_initializer=glorot_uniform(seed=0))(X)
    X = BatchNormalization(axis=4, name=bn_name_base + '2c')(X)

    # Final step: Add shortcut value to main path, and pass it through a RELU activation (≈2 lines)
    X = Add()([X, X_shortcut])
    X = Activation('relu')(X)

    return X


def convolutional_block(X, f, filters, stage, block, s=2):
    """
    Implementation of the convolutional block as defined in Figure 4

    Arguments:
    X -- input tensor of shape (m, n_H_prev, n_W_prev, n_C_prev)
    f -- integer, specifying the shape of the middle CONV's window for the main path
    filters -- python list of integers, defining the number of filters in the CONV layers of the main path
    stage -- integer, used to name the layers, depending on their position in the network
    block -- string/character, used to name the layers, depending on their position in the network
    s -- Integer, specifying the stride to be used

    Returns:
    X -- output of the convolutional block, tensor of shape (n_H, n_W, n_C)
    """

    # defining name basis
    conv_name_base = 'res' + str(stage) + block + '_branch'
    bn_name_base = 'bn' + str(stage) + block + '_branch'

    # Retrieve Filters
    F1, F2, F3 = filters

    # Save the input value
    X_shortcut = X

    ##### MAIN PATH #####
    # First component of main path
    # if stage == 5:
    #     X = Conv3D(filters=F1, kernel_size=(1, 1, 1), strides=(s, s, 1), padding='same', name=conv_name_base + '2a',
    #                kernel_initializer=glorot_uniform(seed=0))(X)
    # else:
    #     X = Conv3D(filters=F1, kernel_size=(1, 1, 1), strides=(s, s, s), padding='same', name=conv_name_base + '2a',
    #                kernel_initializer=glorot_uniform(seed=0))(X)

    X = Conv3D(filters=F1, kernel_size=(1, 1, 1), strides=(s, s, s), padding='same', name=conv_name_base + '2a',
               kernel_initializer=glorot_uniform(seed=0))(X)
    X = BatchNormalization(axis=4, name=bn_name_base + '2a')(X)
    X = Activation('relu')(X)

    # Second component of main path (≈3 lines)
    X = Conv3D(filters=F2, kernel_size=(f, f, f), strides=(1, 1, 1), padding='same', name=conv_name_base + '2b',
               kernel_initializer=glorot_uniform(seed=0))(X)
    X = BatchNormalization(axis=4, name=bn_name_base + '2b')(X)
    X = Activation('relu')(X)

    # Third component of main path (≈2 lines)
    X = Conv3D(filters=F3, kernel_size=(1, 1, 1), strides=(1, 1, 1), padding='same', name=conv_name_base + '2c',
               kernel_initializer=glorot_uniform(seed=0))(X)
    X = BatchNormalization(axis=4, name=bn_name_base + '2c')(X)

    ##### SHORTCUT PATH #### (≈2 lines)
    # if stage == 5:
    #     X_shortcut = Conv3D(filters=F3, kernel_size=(1, 1, 1), strides=(s, s, 1), padding='same',
    #                         name=conv_name_base + '1', kernel_initializer=glorot_uniform(seed=0))(X_shortcut)
    # else:
    #     X_shortcut = Conv3D(filters=F3, kernel_size=(1, 1, 1), strides=(s, s, s), padding='same',
    #                         name=conv_name_base + '1', kernel_initializer=glorot_uniform(seed=0))(X_shortcut)
    X_shortcut = Conv3D(filters=F3, kernel_size=(1, 1, 1), strides=(s, s, s), padding='same',
                        name=conv_name_base + '1', kernel_initializer=glorot_uniform(seed=0))(X_shortcut)
    X_shortcut = BatchNormalization(axis=4, name=bn_name_base + '1')(X_shortcut)

    # Final step: Add shortcut value to main path, and pass it through a RELU activation (≈2 lines)
    X = Add()([X, X_shortcut])
    X = Activation('relu')(X)

    return X


def ResNet50(input_shape=(512, 512, 8, 2), classes=5):
    """
    Implementation of the popular ResNet50 the following architecture:
    CONV2D -> BATCHNORM -> RELU -> MAXPOOL -> CONVBLOCK -> IDBLOCK*2 -> CONVBLOCK -> IDBLOCK*3
    -> CONVBLOCK -> IDBLOCK*5 -> CONVBLOCK -> IDBLOCK*2 -> AVGPOOL -> TOPLAYER

    Arguments:
    input_shape -- shape of the images of the dataset
    classes -- integer, number of classes

    Returns:
    model -- a Model() instance in Keras
    """

    # Define the input as a tensor with shape input_shape
    X_input = Input(input_shape)

    # Stage 1:(512,512,8)-->(256, 256, 4)
    X = Conv3D(64, (3, 3, 3), strides=(1, 1, 1), padding='same', name='conv1', kernel_initializer=glorot_uniform(seed=0))(X_input)
    X = BatchNormalization(axis=4, name='bn_conv1')(X)
    X = Activation('relu')(X)
    X = MaxPooling3D(pool_size=(3, 3, 3), strides=(2, 2, 2))(X)

    # Stage 2:(256, 256, 4)
    X = convolutional_block(X, f=3, filters=[64, 64, 256], stage=2, block='a', s=1)
    X = identity_block(X, 3, [64, 64, 256], stage=2, block='b')
    X = identity_block(X, 3, [64, 64, 256], stage=2, block='c')

    ### START CODE HERE ###

    # Stage 3 (≈4 lines)(256, 256, 4)-->(128, 128, 2)
    X = convolutional_block(X, f=3, filters=[128, 128, 512], stage=3, block='a', s=2)
    X = identity_block(X, 3, [128, 128, 512], stage=3, block='b')
    X = identity_block(X, 3, [128, 128, 512], stage=3, block='c')
    X = identity_block(X, 3, [128, 128, 512], stage=3, block='d')

    # Stage 4 (≈6 lines)(128, 128, 2)-->(64, 64, 1)
    X = convolutional_block(X, f=3, filters=[256, 256, 1024], stage=4, block='a', s=2)
    X = identity_block(X, 3, [256, 256, 1024], stage=4, block='b')
    X = identity_block(X, 3, [256, 256, 1024], stage=4, block='c')
    X = identity_block(X, 3, [256, 256, 1024], stage=4, block='d')
    X = identity_block(X, 3, [256, 256, 1024], stage=4, block='e')
    X = identity_block(X, 3, [256, 256, 1024], stage=4, block='f')

    # Stage 5 (≈3 lines)(64, 64, 1)-->(32,32,1)
    X = convolutional_block(X, f=3, filters=[512, 512, 2048], stage=5, block='a', s=2)
    X = identity_block(X, 3, [512, 512, 2048], stage=5, block='b')
    X = identity_block(X, 3, [512, 512, 2048], stage=5, block='c')

    # AVGPOOL (≈1 line). Use "X = AveragePooling2D(...)(X)"
    X = AveragePooling3D(pool_size=(2, 2, 1), name="avg_pool")(X)

    ### END CODE HERE ###

    # output layer
    X = Flatten()(X)
    X = Dense(128, activation='relu', kernel_initializer=glorot_uniform(seed=0))(X)
    outclass = Dense(classes, activation='softmax', name='outclass')(X)

    # Create model
    model = Model(inputs=X_input, outputs=outclass, name='res3D')
    return model


# 读取二维切片组合为三维输入
def generate_3D_img(path):
    # path = '/data1/zmy/data_augmentation/Slice/3107/CTSlice/'
    ctList = os.listdir(path)
    ctList.sort()

    imglist = []
    for i in range(len(ctList)):
        img = np.load(path+ctList[i])
        imglist.append(img)

    # 对二维图片多于8张的去掉，少的补0
    num = len(imglist)
    if num > 8:
        imglist = imglist[:8]
    elif num < 8:
        for i in range(8-num):
            imglist.append(np.zeros((512, 512)))

    imgs = np.asarray(imglist)

    # z轴变到最后
    imgs = imgs.transpose((1, 2, 0))

    return imgs


# 读取文件列表
def read_csv():
    sets_path = '/data1/zmy/data_augmentation/kfold_dataset/'

    train_sets = ['1.csv',
                  '2.csv',
                  '3.csv',
                  '4.csv']

    test_set = '0.csv'

    # 读取数据集
    train_features = []
    test_features = []

    # 读取训练集
    for set in train_sets:
        train_data = pd.read_csv(sets_path+set)
        for j in range(len(train_data)):
            one_feature = [train_data['patientID'][j], int(train_data['cancer_type'][j])-1]

            train_features.append(one_feature)

    # 读取测试集
    test_data = pd.read_csv(sets_path+test_set)
    for i in range(len(test_data)):
        one_feature = [test_data['patientID'][i], int(test_data['cancer_type'][i])-1]

        test_features.append(one_feature)


    return train_features, test_features


# 训练batch生成器
def data_generator(batch_size, data_list, istrain):
    # 定义返回标志
    batch_idx = 0

    # 样本权重：按类比赋予
    # s_weights = [1, 4, 21, 8, 25]

    # 如果是训练集，打乱数据集
    if istrain:
        random.shuffle(data_list)

    while True:

        # 定义返回特征
        # extra_features = []
        ct_features = []
        pet_features = []

        labels = []
        # sample_weights = []

        for item in data_list:

            labels.append(item[1])
            # sample_weights.append(s_weights[item[15]])

            # extra_features.append(item[:8])

            # 加载三维图像
            ct_path = '/data1/zmy/data_augmentation/Slice/'+str(item[0])+'/CTSlice/'
            ct = generate_3D_img(ct_path)
            ct_features.append(ct)
            pet_path = '/data1/zmy/data_augmentation/Slice/'+str(item[0])+'/PETSlice/'
            pet = generate_3D_img(pet_path)

            # pet图像归一化
            # max_pixel = np.max(pet)
            # min_pixel = np.min(pet)
            # pet = (pet - min_pixel) / float(max_pixel)

            pet_features.append(pet)

            batch_idx = batch_idx+1

            if batch_idx >= batch_size:
                # extra_features = np.asarray(extra_features, dtype=np.float)
                ct_features = np.asarray(ct_features)
                pet_features = np.asarray(pet_features)
                labels = np.asarray(labels, dtype=np.int)

                # 标签转换为onehot编码
                onehot_labels = to_categorical(labels, 5)

                # 特征归一化
                # extra_features = standard_scaler(extra_features)

                # train test ct和pet进行合并
                img_feature = np.stack((ct_features, pet_features), axis=4)
                # print(img_feature.shape)

                # yield {'X_input': img_feature, 'other_info': extra_features}, {'X': onehot_labels}, {'sample_weight': sample_weights}
                yield {'input_1': img_feature}, {'outclass': onehot_labels}

                # extra_features = []
                ct_features = []
                pet_features = []
                labels = []
                # sample_weights = []
                batch_idx = 0


# 加载验证/测试集
def read_testdata(test_features):

    # extra_features = []
    ct_features = []
    pet_features = []
    labels = []
    # sample_weights = []

    for item in test_features:

        labels.append(item[1])
        # sample_weights.append(s_weights[item[15]])

        # extra_features.append(item[:8])

        # 加载三维图像
        ct_path = '/data1/zmy/data_augmentation/Slice/' + str(item[0]) + '/CTSlice/'
        ct = generate_3D_img(ct_path)
        ct_features.append(ct)
        pet_path = '/data1/zmy/data_augmentation/Slice/' + str(item[0]) + '/PETSlice/'
        pet = generate_3D_img(pet_path)

        # pet图像归一化
        # max_pixel = np.max(pet)
        # min_pixel = np.min(pet)
        # pet = (pet - min_pixel) / float(max_pixel)

        pet_features.append(pet)


    # extra_features = np.asarray(extra_features, dtype=np.float)
    ct_features = np.asarray(ct_features)
    pet_features = np.asarray(pet_features)
    labels = np.asarray(labels, dtype=np.int)

    # 标签转换为onehot编码
    onehot_labels = to_categorical(labels, 5)

    # 特征归一化
    # extra_features = standard_scaler(extra_features)

    # train test ct和pet进行合并
    img_feature = np.stack((ct_features, pet_features), axis=4)

    # print(img_feature.shape)

    return img_feature, onehot_labels



def train():
    train_features, test_features = read_csv()

    batch_size = 16

    # 生成训练集和验证集
    train_gen = data_generator(batch_size, train_features, True)

    # 读取验证集
    test_img, test_onehot_labels = read_testdata(test_features)

    model = EasyNet()
    model = multi_gpu_model(model, gpus=2)
    sgd = SGD(lr=0.001, decay=1e-6, momentum=0.9, nesterov=True)
    model.compile(loss='categorical_crossentropy', optimizer=sgd, metrics=['categorical_accuracy'])
    model.summary()

    if not os.path.exists(work_dir):
        os.makedirs(work_dir)

    checkpoint = ModelCheckpoint(filepath=work_dir + "/" + model_name + "_" + "_e" + "{epoch:02d}-{val_loss:.4f}.hd5",
                                 monitor='val_loss',
                                 verbose=1,
                                 save_best_only=False,
                                 save_weights_only=False,
                                 mode='auto',
                                 period=5)


    # 每隔一轮且每当val_loss降低时保存一次模型
    best_model_path = work_dir + "/" + model_name + "_best.hd5"
    checkpoint_fixed_name = ModelCheckpoint(filepath=best_model_path,
                                            monitor='val_loss',
                                            verbose=1,
                                            save_best_only=True,
                                            save_weights_only=False,
                                            mode='auto',
                                            period=1)
    # 不同类别损失权重
    cw = {0: 1, 1: 4, 2: 21, 3: 8, 4: 25}

    # keras 的动态学习率调度，learnrate_scheduler是一个新的学习率
    # learnrate_scheduler = LearningRateScheduler(step_decay)

    model.fit_generator(generator=train_gen,
                        steps_per_epoch=len(train_features) // batch_size,
                        epochs=100,
                        validation_data=(test_img, test_onehot_labels),
                        class_weight=cw,
                        callbacks=[checkpoint, checkpoint_fixed_name, TensorBoard(log_dir=log_dir)])


def predict():
    model_path = "/home/zmy/workspace2021/five_workdirs/five_3D20210315_150702/five_3D__e05-2.1274.hd5"
    model = load_model(model_path)

    train_features, test_features = read_csv()
    test_img, test_onehot_labels = read_testdata(test_features)

    # 测试集推理结果, 这里是五分类概率
    result = model.predict(test_img)

    # 将五分类概率结果转换为下标标签
    preds = np.argmax(result, axis=1)
    labels = np.argmax(test_onehot_labels, axis=1)
    print(preds)
    print(labels)
    # print(preds)
    # print(y_test)

    # 计算评估指标
    target_names = ['1', '2', '3', '4', '5']
    result_statis = classification_report(labels, preds, target_names=target_names)
    print(result_statis)

    # 计算混淆矩阵
    confusion = confusion_matrix(labels, preds)
    print(confusion)


if __name__ == '__main__':
    os.environ["CUDA_VISIBLE_DEVICES"] = "2,3"
    if not os.path.exists(model_save_path):
        os.makedirs(model_save_path)
    # train()
    # best_model_path = work_dir + "/" + model_name + "_best.hd5"
    # shutil.copy(best_model_path, model_save_path)
    predict()

